// Copyright Epic Games, Inc. All Rights Reserved.

#pragma once

#include "zencore.h"

#include <zenbase/concepts.h>

#include <atomic>
#include <string_view>
#include <vector>

namespace zen {
class CbObjectWriter;
}

namespace zen::metrics {

template<typename T>
class Gauge
{
public:
	Gauge() : m_Value{0} {}

	T	 Value() const { return m_Value; }
	void SetValue(T Value) { m_Value = Value; }

private:
	std::atomic<T> m_Value;
};

/** Stats counter
 *
 * A counter is modified by adding or subtracting a value from a current value.
 * This would typically be used to track number of requests in flight, number
 * of active jobs etc
 *
 */
class Counter
{
public:
	inline void		SetValue(uint64_t Value) { m_count = Value; }
	inline uint64_t Value() const { return m_count; }

	inline void Increment(int64_t AddValue) { m_count.fetch_add(AddValue); }
	inline void Decrement(int64_t SubValue) { m_count.fetch_sub(SubValue); }
	inline void Clear() { m_count.store(0, std::memory_order_release); }

private:
	std::atomic<uint64_t> m_count{0};
};

/** Exponential Weighted Moving Average

	This is very raw, to use as little state as possible. If we
	want to use this more broadly in user code we should perhaps
	add a more user-friendly wrapper
 */

class RawEWMA
{
public:
	/// <summary>
	/// Update EWMA with new measure
	/// </summary>
	/// <param name="Alpha">Smoothing factor (between 0 and 1)</param>
	/// <param name="Interval">Elapsed time since last</param>
	/// <param name="Count">Value</param>
	/// <param name="IsInitialUpdate">Whether this is the first update or not</param>
	void   Tick(double Alpha, uint64_t Interval, uint64_t Count, bool IsInitialUpdate);
	double Rate() const;

private:
	std::atomic<double> m_Rate = 0;
};

/// <summary>
/// Tracks rate of events over time (i.e requests/sec), using
/// exponential moving averages
/// </summary>
class Meter
{
public:
	Meter();
	~Meter();

	inline uint64_t Count() const { return m_TotalCount; }
	double			Rate1();				   // One-minute rate
	double			Rate5();				   // Five-minute rate
	double			Rate15();				   // Fifteen-minute rate
	double			MeanRate() const;		   // Mean rate since instantiation of this meter
	void			Mark(uint64_t Count = 1);  // Register one or more events

private:
	std::atomic<uint64_t> m_TotalCount{0};	  // Accumulator counting number of marks since beginning
	std::atomic<uint64_t> m_PendingCount{0};  // Pending EWMA update accumulator
	std::atomic<uint64_t> m_StartTick{0};	  // Time this was instantiated (for mean)
	std::atomic<uint64_t> m_LastTick{0};	  // Timestamp of last EWMA tick
	std::atomic<int64_t>  m_Remainder{0};	  // Tracks the "modulo" of tick time
	bool				  m_IsFirstTick = true;
	RawEWMA				  m_RateM1;
	RawEWMA				  m_RateM5;
	RawEWMA				  m_RateM15;

	void TickIfNecessary();
	void Tick();
};

/** Moment-in-time snapshot of a distribution
 */
class SampleSnapshot
{
public:
	SampleSnapshot(std::vector<double>&& Values);
	~SampleSnapshot();

	uint32_t				   Size() const { return (uint32_t)m_Values.size(); }
	double					   GetQuantileValue(double Quantile);
	double					   GetMedian() { return GetQuantileValue(0.5); }
	double					   Get75Percentile() { return GetQuantileValue(0.75); }
	double					   Get95Percentile() { return GetQuantileValue(0.95); }
	double					   Get98Percentile() { return GetQuantileValue(0.98); }
	double					   Get99Percentile() { return GetQuantileValue(0.99); }
	double					   Get999Percentile() { return GetQuantileValue(0.999); }
	const std::vector<double>& GetValues() const;

private:
	std::vector<double> m_Values;
};

/** Randomly selects samples from a stream. Uses Vitter's
	Algorithm R to produce a statistically representative sample.

  http://www.cs.umd.edu/~samir/498/vitter.pdf - Random Sampling with a Reservoir
 */

class UniformSample
{
public:
	UniformSample(uint32_t ReservoirSize);
	~UniformSample();

	void		   Clear();
	uint32_t	   Size() const;
	void		   Update(int64_t Value);
	SampleSnapshot Snapshot() const;

	template<Invocable<int64_t> T>
	void IterateValues(T Callback) const
	{
		for (const auto& Value : m_Values)
		{
			Callback(Value);
		}
	}

private:
	std::atomic<uint64_t>			  m_SampleCounter{0};
	std::vector<std::atomic<int64_t>> m_Values;
};

/** Track (probabilistic) sample distribution along with min/max
 */
class Histogram
{
public:
	Histogram(int32_t SampleCount = 1028);
	~Histogram();

	void		   Clear();
	void		   Update(int64_t Value);
	int64_t		   Max() const;
	int64_t		   Min() const;
	double		   Mean() const;
	uint64_t	   Count() const;
	SampleSnapshot Snapshot() const { return m_Sample.Snapshot(); }

private:
	UniformSample		 m_Sample;
	std::atomic<int64_t> m_Min{0};
	std::atomic<int64_t> m_Max{0};
	std::atomic<int64_t> m_Sum{0};
	std::atomic<int64_t> m_Count{0};
};

/** Track timing and frequency of some operation

	Example usage would be to track frequency and duration of network
	requests, or function calls.

 */
class OperationTiming
{
public:
	OperationTiming(int32_t SampleCount = 514);
	~OperationTiming();

	void		   Update(int64_t Duration);
	int64_t		   Max() const;
	int64_t		   Min() const;
	double		   Mean() const;
	uint64_t	   Count() const;
	SampleSnapshot Snapshot() const { return m_Histogram.Snapshot(); }

	double Rate1() { return m_Meter.Rate1(); }
	double Rate5() { return m_Meter.Rate5(); }
	double Rate15() { return m_Meter.Rate15(); }
	double MeanRate() const { return m_Meter.MeanRate(); }

	struct Scope
	{
		Scope(OperationTiming& Outer);
		~Scope();

		void Stop();
		void Cancel();

	private:
		OperationTiming& m_Outer;
		uint64_t		 m_StartTick;
	};

private:
	Meter	  m_Meter;
	Histogram m_Histogram;
};

struct MeterSnapshot
{
	uint64_t Count;
	double	 MeanRate;
	double	 Rate1;
	double	 Rate5;
	double	 Rate15;
};

struct HistogramSnapshot
{
	double Count;
	double Avg;
	double Min;
	double Max;
	double P75;
	double P95;
	double P99;
	double P999;
};

struct StatsSnapshot
{
	MeterSnapshot	  Meter;
	HistogramSnapshot Histogram;
};

struct RequestStatsSnapshot
{
	StatsSnapshot Requests;
	StatsSnapshot Bytes;
};

/** Metrics for network requests

	Aggregates tracking of duration, payload sizes into a single
	class

  */
class RequestStats
{
public:
	RequestStats(int32_t SampleCount = 514);
	~RequestStats();

	void	 Update(int64_t Duration, int64_t Bytes);
	uint64_t Count() const;

	// Timing

	int64_t		   MaxDuration() const { return m_BytesHistogram.Max(); }
	int64_t		   MinDuration() const { return m_BytesHistogram.Min(); }
	double		   MeanDuration() const { return m_BytesHistogram.Mean(); }
	SampleSnapshot DurationSnapshot() const { return m_RequestTimeHistogram.Snapshot(); }
	double		   Rate1() { return m_RequestMeter.Rate1(); }
	double		   Rate5() { return m_RequestMeter.Rate5(); }
	double		   Rate15() { return m_RequestMeter.Rate15(); }
	double		   MeanRate() const { return m_RequestMeter.MeanRate(); }

	// Bytes

	int64_t		   MaxBytes() const { return m_BytesHistogram.Max(); }
	int64_t		   MinBytes() const { return m_BytesHistogram.Min(); }
	double		   MeanBytes() const { return m_BytesHistogram.Mean(); }
	SampleSnapshot BytesSnapshot() const { return m_BytesHistogram.Snapshot(); }
	double		   ByteRate1() { return m_BytesMeter.Rate1(); }
	double		   ByteRate5() { return m_BytesMeter.Rate5(); }
	double		   ByteRate15() { return m_BytesMeter.Rate15(); }
	double		   ByteMeanRate() const { return m_BytesMeter.MeanRate(); }

	struct Scope
	{
		Scope(RequestStats& Outer, int64_t Bytes);
		~Scope();

		void SetBytes(int64_t Bytes) { m_Bytes = Bytes; }
		void Stop();
		void Cancel();

	private:
		RequestStats& m_Outer;
		uint64_t	  m_StartTick;
		int64_t		  m_Bytes;
	};

	void EmitSnapshot(std::string_view Tag, CbObjectWriter& Cbo);

	RequestStatsSnapshot Snapshot();

private:
	static StatsSnapshot GetSnapshot(Meter& M, Histogram& H, double ConversionFactor)
	{
		SampleSnapshot Snap = H.Snapshot();
		return StatsSnapshot{
			.Meter	   = {.Count = M.Count(), .MeanRate = M.MeanRate(), .Rate1 = M.Rate1(), .Rate5 = M.Rate5(), .Rate15 = M.Rate15()},
			.Histogram = {.Count = H.Count() * ConversionFactor,
						  .Avg	 = H.Mean() * ConversionFactor,
						  .Min	 = H.Min() * ConversionFactor,
						  .Max	 = H.Max() * ConversionFactor,
						  .P75	 = Snap.Get75Percentile() * ConversionFactor,
						  .P95	 = Snap.Get95Percentile() * ConversionFactor,
						  .P99	 = Snap.Get99Percentile() * ConversionFactor,
						  .P999	 = Snap.Get999Percentile() * ConversionFactor}};
	}

	Meter	  m_RequestMeter;
	Meter	  m_BytesMeter;
	Histogram m_RequestTimeHistogram;
	Histogram m_BytesHistogram;
};

void EmitSnapshot(std::string_view Tag, OperationTiming& Stat, CbObjectWriter& Cbo);
void EmitSnapshot(std::string_view Tag, const Histogram& Stat, CbObjectWriter& Cbo, double ConversionFactor);
void EmitSnapshot(std::string_view Tag, Meter& Stat, CbObjectWriter& Cbo);

void EmitSnapshot(const Histogram& Stat, CbObjectWriter& Cbo, double ConversionFactor);

void EmitSnapshot(std::string_view Tag, const MeterSnapshot& Snapshot, CbObjectWriter& Cbo);
void EmitSnapshot(std::string_view Tag, const HistogramSnapshot& Snapshot, CbObjectWriter& Cbo);
void EmitSnapshot(std::string_view Tag, const StatsSnapshot& Snapshot, CbObjectWriter& Cbo);
void EmitSnapshot(std::string_view Tag, const RequestStatsSnapshot& Snapshot, CbObjectWriter& Cbo);

}  // namespace zen::metrics

namespace zen {

extern void stats_forcelink();

}  // namespace zen
